from LagrangianParticles import circle
from dolfin import *
import numpy
from numpy import sin,cos
from copy import deepcopy


'''Enright test.'''

def build(points):
    '''Given ordered points build boundary mesh and the level-set function.''' 
    N = len(points)
    
    mesh = Mesh()
    editor = MeshEditor()
    editor.open(mesh, 1, 2)
    editor.init_vertices(N)
    editor.init_cells(N)

    for i, (x, y) in enumerate(points):
        editor.add_vertex(i, x, y)
    for i in range(N):
        editor.add_cell(i, i, (i+1) % (N))
    editor.close()

    rect = RectangleMesh(0,0,1,1,50,50)
    V = FunctionSpace(rect, 'CG', 1)
    x0 = interpolate(Expression('x[0]'), V)
    x1 = interpolate(Expression('x[1]'), V)
    a0 = x0.vector().array()
    a1 = x1.vector().array()

    # For all nodes in V find the signed distance to the disc 
    point_outside = Point(0,0)  # A point outside the disc
    Np = x0.vector().size()
    dummymesh = Mesh()
    dummyeditor = MeshEditor()
    dummyeditor.open(dummymesh, 1, 2)
    dummyeditor.init_vertices(Np + 1)
    dummyeditor.init_cells(Np)
    dummyeditor.add_vertex(0, point_outside.x(), point_outside.y())

    inside = Function(V)
    inside.vector()[:] = 1
    # Use ray casting to determine if the point is inside (-1) or outside (+1)
    # There's probably a better way to do this than to create a dummy mesh
    for i in range(Np):
        dummyeditor.add_vertex(i + 1, a0[i], a1[i])
        dummyeditor.add_cell(i, 0, i + 1)
        ei = Edge(dummymesh, i) # Edge (ray) from vertex 0 to i + 1
        ix = mesh.intersected_cells(ei) # Odd number means you're inside
        if ix.shape[0] % 2 == 1:
            inside.vector()[i] = -1
    dummyeditor.close() 
    phiX_ = Function(V)

    i = 0
    # Compute the signed distance
    for xx, yy in zip(a0, a1):
        p = Point(xx, yy)
        zz = mesh.distance(p)
        phiX_.vector()[i] = zz*inside.vector()[i]
        i += 1
    
    # compute length of the mesh
    mesh.order()
    # plot(mesh,interactive=True)
    W = FunctionSpace(mesh,'CG',1)
    L = assemble(Constant(1.0)*dx,mesh=mesh)

    return mesh,phiX_,L

def move(points,dt):
    '''3rd order TVD time stepping.'''
    points0 = deepcopy(points) # initial position of particles

    # stage one
    for point in points:
        x = deepcopy(point[0])
        y = deepcopy(point[1])
        A = numpy.array([-2*sin(pi*y)*cos(pi*y)*sin(pi*x)*sin(pi*x),2*sin(pi*x)*cos(pi*x)*sin(pi*y)*sin(pi*y)])
        
        point[:] = point[:] + dt*A

    # stage two
    for i,point in enumerate(points):
        x = deepcopy(point[0])
        y = deepcopy(point[1])

        A = numpy.array([-2*sin(pi*y)*cos(pi*y)*sin(pi*x)*sin(pi*x),2*sin(pi*x)*cos(pi*x)*sin(pi*y)*sin(pi*y)])
        
        point[:] = 0.75*points0[i][:] + 0.25*(point[:] + dt*A)

    # stage three
    for i,point in enumerate(points):
        x = deepcopy(point[0])
        y = deepcopy(point[1])

        A = numpy.array([-2*sin(pi*y)*cos(pi*y)*sin(pi*x)*sin(pi*x),2*sin(pi*x)*cos(pi*x)*sin(pi*y)*sin(pi*y)])
        
        point[:] = (1./3.)*points0[i][:] + (2./3.)*(point[:] + dt*A)



points = circle(x0=(0.5,0.75),N=1000)

t = 0
dt = 1E-3
step = 0
out = File('enright.pvd')

count = 1
while t < 1:
    move(points,dt)
    bdry_mesh, sdf, L = build(points)
    t += dt
    if step%100 == 0:
        out << sdf, count
        #plot(bdry_mesh)
        #plot(sdf)
        #interactive()

        print t, L
        count += 1

    step += 1


